# -*- encoding: utf-8 -*-
'''
@File    :   Middleware.py   
@Contact :   desufnocs@163.com
@des     :   中间件
 
@Modify Time       @Author      @Version    @Desciption
------------      ----------    --------    -----------
2025/2/18 14:57   AssenImpact      1.0         None
'''
import re

from fastapi import status,HTTPException
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from control import config as CONFIG
from control.core.Response import fail
from control.core.Utils import verify_token


class IPMiddleware:
    def __init__(self, app: ASGIApp, allow_origins: list[str] | None = None) -> None:
        self.app = app
        self.allow_origins = allow_origins or ["*"]

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        fail_response = JSONResponse(
            status_code=status.HTTP_502_BAD_GATEWAY,
            content=fail(code=502, msg="Nothing"),
        )
        if scope["type"] == "http":
            real_ip = scope["client"][0]
            access = False
            for rx in self.allow_origins:
                if re.fullmatch(rx, real_ip):
                    access = True
            if not access:
                return await fail_response(scope, receive, send)
        await self.app(scope, receive, send)


class VerifyTokenMiddleware:
    def __init__(self, app: ASGIApp, access_header: dict[str] | None = None) -> None:
        self.app = app
        self.access_header = access_header or {}

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        fail_response = JSONResponse(
            status_code=status.HTTP_403_FORBIDDEN,
            content=fail(code=403, msg="who are you?"),
        )
        if scope["type"] == "http":  # 仅处理 HTTP 请求
            headers = {k.decode(): v.decode() for k, v in scope["headers"]}
            if not CONFIG.DEVELOPER_MODE:
                if scope["method"] == "OPTIONS":
                    access_request_header = set(headers.get("access-control-request-headers", "").split(","))
                    if not access_request_header.issuperset(self.access_header):
                        return await fail_response(scope, receive, send)
                else:
                    # 检查 X-Access-Key 头
                    token = headers.get("x-access-key")
                    try:
                        if not token or not verify_token(token):
                            return await fail_response(scope, receive, send)
                    except:
                        return await fail_response(scope, receive, send)
                    # 继续处理请求
        await self.app(scope, receive, send)
